package vectordb;

import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.time.Duration;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;

/**
 * 基于 Milvus RESTful API 的向量数据库实现
 */
public class MilvusRestfulVectorDB implements vectordb.VectorDB {
    private String host;
    private int port;
    private String username;
    private String password;
    private boolean secure;
    private String baseUrl;
    private boolean initialized;
    private HttpClient client;
    private ObjectMapper objectMapper;
    private String token;

    /**
     * Milvus RESTful 配置类
     */
    public static class MilvusRestfulConfig {
        private String host;
        private int port;
        private String username;
        private String password;
        private boolean secure;

        public MilvusRestfulConfig(String host, int port) {
            this.host = host;
            this.port = port;
            this.secure = false;
        }

        public MilvusRestfulConfig(String host, int port, String username, String password) {
            this.host = host;
            this.port = port;
            this.username = username;
            this.password = password;
            this.secure = false;
        }

        public MilvusRestfulConfig(String host, int port, String username, String password, boolean secure) {
            this.host = host;
            this.port = port;
            this.username = username;
            this.password = password;
            this.secure = secure;
        }

        public String getHost() {
            return host;
        }

        public void setHost(String host) {
            this.host = host;
        }

        public int getPort() {
            return port;
        }

        public void setPort(int port) {
            this.port = port;
        }

        public String getUsername() {
            return username;
        }

        public void setUsername(String username) {
            this.username = username;
        }

        public String getPassword() {
            return password;
        }

        public void setPassword(String password) {
            this.password = password;
        }

        public boolean isSecure() {
            return secure;
        }

        public void setSecure(boolean secure) {
            this.secure = secure;
        }
    }

    /**
     * 创建一个新的 Milvus RESTful 向量数据库连接
     * 
     * @param config Milvus RESTful 配置
     */
    public MilvusRestfulVectorDB(MilvusRestfulConfig config) {
        this.host = config.getHost();
        this.port = config.getPort();
        this.username = config.getUsername();
        this.password = config.getPassword();
        this.secure = config.isSecure();
        
        String protocol = this.secure ? "https" : "http";
        this.baseUrl = protocol + "://" + this.host + ":" + this.port + "/api/v1";
        
        this.initialized = false;
        this.client = HttpClient.newBuilder()
                .connectTimeout(Duration.ofSeconds(30))
                .build();
        this.objectMapper = new ObjectMapper();
    }

    @Override
    public CompletableFuture<Void> init() {
        if (this.initialized) {
            return CompletableFuture.completedFuture(null);
        }
        
        return CompletableFuture.supplyAsync(() -> {
            try {
                // 如果提供了用户名和密码，则进行身份验证
                if (this.username != null && !this.username.isEmpty() && this.password != null) {
                    Map<String, String> authBody = new HashMap<>();
                    authBody.put("username", this.username);
                    authBody.put("password", this.password);
                    
                    String authBodyJson = objectMapper.writeValueAsString(authBody);
                    
                    HttpRequest authRequest = HttpRequest.newBuilder()
                            .uri(URI.create(this.baseUrl + "/user/login"))
                            .header("Content-Type", "application/json")
                            .POST(HttpRequest.BodyPublishers.ofString(authBodyJson))
                            .build();
                    
                    HttpResponse<String> authResponse = client.send(authRequest, HttpResponse.BodyHandlers.ofString());
                    
                    if (authResponse.statusCode() != 200) {
                        throw new RuntimeException("Failed to authenticate with Milvus: " + authResponse.body());
                    }
                    
                    Map<String, Object> authResponseMap = objectMapper.readValue(authResponse.body(), Map.class);
                    Map<String, Object> data = (Map<String, Object>) authResponseMap.get("data");
                    this.token = (String) data.get("token");
                }
                
                this.initialized = true;
                System.out.println("Connected to Milvus RESTful API at " + this.baseUrl);
                return null;
            } catch (Exception e) {
                System.err.println("Failed to connect to Milvus RESTful API: " + e.getMessage());
                throw new RuntimeException("Failed to connect to Milvus RESTful API", e);
            }
        });
    }

    @Override
    public CompletableFuture<Void> createCollection(String collectionName, int dimension, Map<String, String> metadataFields) {
        return this.init()
                .thenCompose(result -> CompletableFuture.supplyAsync(() -> {
                    try {
                        Map<String, Object> requestBody = new HashMap<>();
                        requestBody.put("collection_name", collectionName);
                        
                        // 构建字段定义
                        List<Map<String, Object>> fields = new ArrayList<>();
                        
                        // ID 字段
                        Map<String, Object> idField = new HashMap<>();
                        idField.put("name", "id");
                        idField.put("description", "ID field");
                        idField.put("data_type", "VarChar");
                        idField.put("is_primary_key", true);
                        idField.put("max_length", 100);
                        fields.add(idField);
                        
                        // 向量字段
                        Map<String, Object> vectorField = new HashMap<>();
                        vectorField.put("name", "vector");
                        vectorField.put("description", "Vector field");
                        vectorField.put("data_type", "FloatVector");
                        vectorField.put("dim", dimension);
                        fields.add(vectorField);
                        
                        // 内容字段
                        Map<String, Object> contentField = new HashMap<>();
                        contentField.put("name", "content");
                        contentField.put("description", "Content field");
                        contentField.put("data_type", "VarChar");
                        contentField.put("max_length", 65535);
                        fields.add(contentField);
                        
                        // 路径字段
                        Map<String, Object> pathField = new HashMap<>();
                        pathField.put("name", "relativePath");
                        pathField.put("description", "Relative path field");
                        pathField.put("data_type", "VarChar");
                        pathField.put("max_length", 1024);
                        fields.add(pathField);
                        
                        // 行范围字段
                        Map<String, Object> startLineField = new HashMap<>();
                        startLineField.put("name", "startLine");
                        startLineField.put("description", "Start line field");
                        startLineField.put("data_type", "Int64");
                        fields.add(startLineField);
                        
                        Map<String, Object> endLineField = new HashMap<>();
                        endLineField.put("name", "endLine");
                        endLineField.put("description", "End line field");
                        endLineField.put("data_type", "Int64");
                        fields.add(endLineField);
                        
                        // 文件扩展名字段
                        Map<String, Object> fileExtField = new HashMap<>();
                        fileExtField.put("name", "fileExtension");
                        fileExtField.put("description", "File extension field");
                        fileExtField.put("data_type", "VarChar");
                        fileExtField.put("max_length", 50);
                        fields.add(fileExtField);
                        
                        // 添加元数据字段
                        if (metadataFields != null) {
                            for (Map.Entry<String, String> entry : metadataFields.entrySet()) {
                                Map<String, Object> metaField = new HashMap<>();
                                metaField.put("name", entry.getKey());
                                metaField.put("description", "Metadata field: " + entry.getKey());
                                
                                String dataType = entry.getValue();
                                metaField.put("data_type", dataType);
                                
                                if (dataType.equals("VarChar")) {
                                    metaField.put("max_length", 1024);
                                }
                                
                                fields.add(metaField);
                            }
                        }
                        
                        requestBody.put("fields", fields);
                        
                        // 索引参数
                        Map<String, Object> indexParams = new HashMap<>();
                        indexParams.put("index_type", "HNSW");
                        indexParams.put("metric_type", "L2");
                        indexParams.put("params", Map.of("M", 8, "efConstruction", 64));
                        
                        requestBody.put("index_params", indexParams);
                        
                        String requestBodyJson = objectMapper.writeValueAsString(requestBody);
                        
                        HttpRequest.Builder requestBuilder = HttpRequest.newBuilder()
                                .uri(URI.create(this.baseUrl + "/collections"))
                                .header("Content-Type", "application/json");
                        
                        if (this.token != null) {
                            requestBuilder.header("Authorization", "Bearer " + this.token);
                        }
                        
                        HttpRequest request = requestBuilder
                                .POST(HttpRequest.BodyPublishers.ofString(requestBodyJson))
                                .build();
                        
                        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
                        
                        if (response.statusCode() != 200) {
                            throw new RuntimeException("Failed to create collection: " + response.body());
                        }
                        
                        System.out.println("Created collection: " + collectionName);
                        return null;
                    } catch (Exception e) {
                        System.err.println("Failed to create collection: " + e.getMessage());
                        throw new RuntimeException("Failed to create collection", e);
                    }
                }));
    }

    @Override
    public CompletableFuture<Boolean> hasCollection(String collectionName) {
        return this.init()
                .thenCompose(result -> CompletableFuture.supplyAsync(() -> {
                    try {
                        HttpRequest.Builder requestBuilder = HttpRequest.newBuilder()
                                .uri(URI.create(this.baseUrl + "/collections/" + collectionName));
                        
                        if (this.token != null) {
                            requestBuilder.header("Authorization", "Bearer " + this.token);
                        }
                        
                        HttpRequest request = requestBuilder.GET().build();
                        
                        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
                        
                        return response.statusCode() == 200;
                    } catch (Exception e) {
                        System.err.println("Failed to check collection: " + e.getMessage());
                        return false;
                    }
                }));
    }

    @Override
    public CompletableFuture<Void> dropCollection(String collectionName) {
        return this.init()
                .thenCompose(result -> CompletableFuture.supplyAsync(() -> {
                    try {
                        HttpRequest.Builder requestBuilder = HttpRequest.newBuilder()
                                .uri(URI.create(this.baseUrl + "/collections/" + collectionName));
                        
                        if (this.token != null) {
                            requestBuilder.header("Authorization", "Bearer " + this.token);
                        }
                        
                        HttpRequest request = requestBuilder.DELETE().build();
                        
                        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
                        
                        if (response.statusCode() != 200) {
                            throw new RuntimeException("Failed to drop collection: " + response.body());
                        }
                        
                        System.out.println("Dropped collection: " + collectionName);
                        return null;
                    } catch (Exception e) {
                        System.err.println("Failed to drop collection: " + e.getMessage());
                        throw new RuntimeException("Failed to drop collection", e);
                    }
                }));
    }

    @Override
    public CompletableFuture<Void> insert(String collectionName, List<vectordb.VectorDocument> documents) {
        return this.init()
                .thenCompose(result -> CompletableFuture.supplyAsync(() -> {
                    try {
                        // 确保每个文档都有 ID
                        documents.forEach(doc -> {
                            if (doc.getId() == null || doc.getId().isEmpty()) {
                                doc.setId(UUID.randomUUID().toString());
                            }
                        });
                        
                        Map<String, Object> requestBody = new HashMap<>();
                        requestBody.put("collection_name", collectionName);
                        
                        // 构建数据
                        List<Map<String, Object>> data = new ArrayList<>();
                        for (vectordb.VectorDocument doc : documents) {
                            Map<String, Object> entity = new HashMap<>();
                            entity.put("id", doc.getId());
                            entity.put("vector", doc.getVector());
                            entity.put("content", doc.getContent());
                            entity.put("relativePath", doc.getRelativePath());
                            entity.put("startLine", doc.getStartLine());
                            entity.put("endLine", doc.getEndLine());
                            entity.put("fileExtension", doc.getFileExtension());
                            
                            // 添加元数据
                            if (doc.getMetadata() != null) {
                                for (Map.Entry<String, Object> entry : doc.getMetadata().entrySet()) {
                                    entity.put(entry.getKey(), entry.getValue());
                                }
                            }
                            
                            data.add(entity);
                        }
                        
                        requestBody.put("data", data);
                        
                        String requestBodyJson = objectMapper.writeValueAsString(requestBody);
                        
                        HttpRequest.Builder requestBuilder = HttpRequest.newBuilder()
                                .uri(URI.create(this.baseUrl + "/entities"))
                                .header("Content-Type", "application/json");
                        
                        if (this.token != null) {
                            requestBuilder.header("Authorization", "Bearer " + this.token);
                        }
                        
                        HttpRequest request = requestBuilder
                                .POST(HttpRequest.BodyPublishers.ofString(requestBodyJson))
                                .build();
                        
                        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
                        
                        if (response.statusCode() != 200) {
                            throw new RuntimeException("Failed to insert documents: " + response.body());
                        }
                        
                        System.out.println("Inserted " + documents.size() + " documents into collection: " + collectionName);
                        return null;
                    } catch (Exception e) {
                        System.err.println("Failed to insert documents: " + e.getMessage());
                        throw new RuntimeException("Failed to insert documents", e);
                    }
                }));
    }

    @Override
    public CompletableFuture<List<vectordb.VectorSearchResult>> search(String collectionName, double[] queryVector, int limit, String filter) {
        return this.init()
                .thenCompose(result -> CompletableFuture.supplyAsync(() -> {
                    try {
                        Map<String, Object> requestBody = new HashMap<>();
                        requestBody.put("collection_name", collectionName);
                        requestBody.put("vector", queryVector);
                        requestBody.put("limit", limit);
                        
                        // 添加过滤条件
                        if (filter != null && !filter.isEmpty()) {
                            requestBody.put("filter", filter);
                        }
                        
                        // 指定输出字段
                        List<String> outputFields = List.of("id", "content", "relativePath", "startLine", "endLine", "fileExtension");
                        requestBody.put("output_fields", outputFields);
                        
                        String requestBodyJson = objectMapper.writeValueAsString(requestBody);
                        
                        HttpRequest.Builder requestBuilder = HttpRequest.newBuilder()
                                .uri(URI.create(this.baseUrl + "/search"))
                                .header("Content-Type", "application/json");
                        
                        if (this.token != null) {
                            requestBuilder.header("Authorization", "Bearer " + this.token);
                        }
                        
                        HttpRequest request = requestBuilder
                                .POST(HttpRequest.BodyPublishers.ofString(requestBodyJson))
                                .build();
                        
                        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
                        
                        if (response.statusCode() != 200) {
                            throw new RuntimeException("Failed to search: " + response.body());
                        }
                        
                        Map<String, Object> responseMap = objectMapper.readValue(response.body(), Map.class);
                        List<Map<String, Object>> results = (List<Map<String, Object>>) responseMap.get("results");
                        
                        List<vectordb.VectorSearchResult> searchResults = new ArrayList<>();
                        for (Map<String, Object> result : results) {
                            String id = (String) result.get("id");
                            String content = (String) result.get("content");
                            String relativePath = (String) result.get("relativePath");
                            int startLine = ((Number) result.get("startLine")).intValue();
                            int endLine = ((Number) result.get("endLine")).intValue();
                            String fileExtension = (String) result.get("fileExtension");
                            double score = ((Number) result.get("score")).doubleValue();
                            
                            vectordb.VectorDocument doc = new vectordb.VectorDocument(
                                id, null, content, relativePath, startLine, endLine, fileExtension
                            );
                            
                            searchResults.add(new vectordb.VectorSearchResult(doc, score));
                        }
                        
                        return searchResults;
                    } catch (Exception e) {
                        System.err.println("Failed to search: " + e.getMessage());
                        return new ArrayList<>();
                    }
                }));
    }

    @Override
    public CompletableFuture<Long> count(String collectionName) {
        return this.init()
                .thenCompose(result -> CompletableFuture.supplyAsync(() -> {
                    try {
                        HttpRequest.Builder requestBuilder = HttpRequest.newBuilder()
                                .uri(URI.create(this.baseUrl + "/collections/" + collectionName + "/stats"));
                        
                        if (this.token != null) {
                            requestBuilder.header("Authorization", "Bearer " + this.token);
                        }
                        
                        HttpRequest request = requestBuilder.GET().build();
                        
                        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
                        
                        if (response.statusCode() != 200) {
                            throw new RuntimeException("Failed to get collection stats: " + response.body());
                        }
                        
                        Map<String, Object> responseMap = objectMapper.readValue(response.body(), Map.class);
                        Map<String, Object> data = (Map<String, Object>) responseMap.get("data");
                        Map<String, Object> stats = (Map<String, Object>) data.get("stats");
                        
                        return ((Number) stats.get("row_count")).longValue();
                    } catch (Exception e) {
                        System.err.println("Failed to count documents: " + e.getMessage());
                        return 0L;
                    }
                }));
    }

    @Override
    public CompletableFuture<Void> delete(String collectionName, List<String> ids) {
        return this.init()
                .thenCompose(result -> CompletableFuture.supplyAsync(() -> {
                    try {
                        Map<String, Object> requestBody = new HashMap<>();
                        requestBody.put("collection_name", collectionName);
                        requestBody.put("ids", ids);
                        
                        String requestBodyJson = objectMapper.writeValueAsString(requestBody);
                        
                        HttpRequest.Builder requestBuilder = HttpRequest.newBuilder()
                                .uri(URI.create(this.baseUrl + "/entities"))
                                .header("Content-Type", "application/json");
                        
                        if (this.token != null) {
                            requestBuilder.header("Authorization", "Bearer " + this.token);
                        }
                        
                        HttpRequest request = requestBuilder
                                .method("DELETE", HttpRequest.BodyPublishers.ofString(requestBodyJson))
                                .build();
                        
                        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
                        
                        if (response.statusCode() != 200) {
                            throw new RuntimeException("Failed to delete documents: " + response.body());
                        }
                        
                        System.out.println("Deleted " + ids.size() + " documents from collection: " + collectionName);
                        return null;
                    } catch (Exception e) {
                        System.err.println("Failed to delete documents: " + e.getMessage());
                        throw new RuntimeException("Failed to delete documents", e);
                    }
                }));
    }

    @Override
    public CompletableFuture<Void> close() {
        // RESTful API 不需要关闭连接
        this.initialized = false;
        return CompletableFuture.completedFuture(null);
    }
}